package com.liyuxiang.graph.config;

import com.alibaba.cloud.ai.graph.GraphRepresentation;
import com.alibaba.cloud.ai.graph.KeyStrategy;
import com.alibaba.cloud.ai.graph.KeyStrategyFactory;
import com.alibaba.cloud.ai.graph.StateGraph;
import com.alibaba.cloud.ai.graph.action.AsyncNodeAction;
import com.alibaba.cloud.ai.graph.exception.GraphStateException;
import com.alibaba.cloud.ai.graph.state.strategy.ReplaceStrategy;
import com.alibaba.cloud.ai.graph.state.strategy.AppendStrategy;
import com.liyuxiang.graph.node.*;
import com.liyuxiang.graph.service.*;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.CompletableFuture;


@Configuration
public class GraphConfiguration {

    @Bean("diagnosisGraph")
    public StateGraph diagnosisGraph(ChatClient.Builder chatClientBuilder,
                                     IntentRecognitionService intentRecognitionService,
                                     SymptomAnalysisService symptomAnalysisService,
                                     DepartmentRecommendationService departmentRecommendationService,
                                     ConversationHandlingService conversationHandlingService,
                                     DiagnosisContextService contextService) throws GraphStateException {

        // 全局变量的替换策略
        KeyStrategyFactory keyStrategyFactory = () -> {
            HashMap<String, KeyStrategy> keyStrategyHashMap = new HashMap<>();
            // 用户输入和状态管理
            keyStrategyHashMap.put("query", new ReplaceStrategy());
            keyStrategyHashMap.put("userId", new ReplaceStrategy());
            keyStrategyHashMap.put("intentType", new ReplaceStrategy());
            keyStrategyHashMap.put("symptoms", new ReplaceStrategy());
            keyStrategyHashMap.put("conversationHistory", new AppendStrategy());
            keyStrategyHashMap.put("questionCount", new ReplaceStrategy());
            keyStrategyHashMap.put("recommendedDepartment", new ReplaceStrategy());
            keyStrategyHashMap.put("response", new ReplaceStrategy());
            keyStrategyHashMap.put("isCompleted", new ReplaceStrategy());
            keyStrategyHashMap.put("shouldContinue", new ReplaceStrategy());
            return keyStrategyHashMap;
        };

        // 构造智能问诊图
        StateGraph diagnosisGraph = new StateGraph(keyStrategyFactory)
                // 意图识别节点
                .addNode("intentRecognition", AsyncNodeAction.node_async(new IntentRecognitionNode(intentRecognitionService, contextService)))
                // 症状详情收集节点
                .addNode("symptomDetail", AsyncNodeAction.node_async(new SymptomDetailNode(symptomAnalysisService)))
                // 科室推荐节点
                .addNode("departmentRecommendation", AsyncNodeAction.node_async(new DepartmentRecommendationNode(departmentRecommendationService)))
                // 无意图问题处理节点
                .addNode("noIntentQuestion", AsyncNodeAction.node_async(new NoIntentQuestionNode(conversationHandlingService)))

                // 边：START -> 意图识别
                .addEdge(StateGraph.START, "intentRecognition")

                // 条件边：意图识别 -> 根据意图类型分流
                .addConditionalEdges("intentRecognition",
                        state -> {
                            String intentType = state.value("intentType", "无意图");
                            System.out.println("=== 意图识别结果: " + intentType + " ===");
                            if ("有意图".equals(intentType)) {
                                return CompletableFuture.completedFuture("symptomDetail");
                            } else if ("直接挂号".equals(intentType)) {
                                System.out.println("=== 检测到直接挂号，跳转到科室推荐 ===");
                                return CompletableFuture.completedFuture("departmentRecommendation");
                            } else {
                                return CompletableFuture.completedFuture("noIntentQuestion");
                            }
                        },
                        Map.of(
                                "symptomDetail", "symptomDetail",
                                "departmentRecommendation", "departmentRecommendation",
                                "noIntentQuestion", "noIntentQuestion"
                        ))

                // 条件边：症状详情收集 -> 判断是否继续询问或推荐科室
                .addConditionalEdges("symptomDetail",
                        state -> {
                            Boolean shouldContinue = state.value("shouldContinue", true);
                            Integer questionCount = state.value("questionCount", 0);
                            String suggestedAction = state.value("suggestedAction", "continue");
                            System.out.println("=== 症状收集状态 - 问题轮次: " + questionCount + ", 是否继续: " + shouldContinue + ", 建议行动: " + suggestedAction + " ===");

                            // 如果建议重定向到对话处理
                            if ("redirect_to_conversation".equals(suggestedAction)) {
                                return CompletableFuture.completedFuture("noIntentQuestion");
                            }

                            if (shouldContinue && questionCount < 5) {
                                return CompletableFuture.completedFuture(StateGraph.END); // 暂时结束，等待用户回复
                            } else {
                                return CompletableFuture.completedFuture("departmentRecommendation");
                            }
                        },
                        Map.of(
                                StateGraph.END, StateGraph.END,
                                "departmentRecommendation", "departmentRecommendation",
                                "noIntentQuestion", "noIntentQuestion"
                        ))

                // 边：科室推荐 -> END
                .addEdge("departmentRecommendation", StateGraph.END)
                // 边：无意图问题 -> END
                .addEdge("noIntentQuestion", StateGraph.END);

        // 打印图结构
        GraphRepresentation representation = diagnosisGraph.getGraph(GraphRepresentation.Type.PLANTUML,
                "智能问诊流程图");
        System.out.println("\n=== 智能问诊 UML Flow ===");
        System.out.println(representation.content());
        System.out.println("==================================\n");

        return diagnosisGraph;
    }

    /**
     * 流式响应的智能问诊图配置
     */
    @Bean("streamDiagnosisGraph")
    public StateGraph streamDiagnosisGraph(ChatClient.Builder chatClientBuilder,
                                           DiagnosisContextService contextService,
                                           IntentRecognitionService intentRecognitionService,
                                           SymptomAnalysisService symptomAnalysisService,
                                           DepartmentRecommendationService departmentRecommendationService,
                                           ConversationHandlingService conversationHandlingService) throws GraphStateException {

        // 流式图的键策略配置
        KeyStrategyFactory keyStrategyFactory = () -> {
            HashMap<String, KeyStrategy> keyStrategyHashMap = new HashMap<>();
            // 基础状态管理
            keyStrategyHashMap.put("query", new ReplaceStrategy());
            keyStrategyHashMap.put("userId", new ReplaceStrategy());
            keyStrategyHashMap.put("threadId", new ReplaceStrategy());
            keyStrategyHashMap.put("intentType", new ReplaceStrategy());
            keyStrategyHashMap.put("symptoms", new ReplaceStrategy());
            keyStrategyHashMap.put("conversationHistory", new AppendStrategy());
            keyStrategyHashMap.put("questionCount", new ReplaceStrategy());
            keyStrategyHashMap.put("recommendedDepartment", new ReplaceStrategy());
            keyStrategyHashMap.put("response", new ReplaceStrategy());
            keyStrategyHashMap.put("isCompleted", new ReplaceStrategy());
            keyStrategyHashMap.put("shouldContinue", new ReplaceStrategy());
            keyStrategyHashMap.put("readyForRecommendation", new ReplaceStrategy());
            keyStrategyHashMap.put("completenessScore", new ReplaceStrategy());
            keyStrategyHashMap.put("reasoning", new ReplaceStrategy());

            // 流式输出相关
            keyStrategyHashMap.put("intentAnalysis", new ReplaceStrategy());
            keyStrategyHashMap.put("symptomDetailAnalysis", new ReplaceStrategy());
            keyStrategyHashMap.put("departmentAnalysis", new ReplaceStrategy());
            keyStrategyHashMap.put("conversationAnalysis", new ReplaceStrategy());

            return keyStrategyHashMap;
        };

        // 构建流式智能问诊图
        StateGraph streamGraph = new StateGraph(keyStrategyFactory)
                // 流式意图识别节点
                .addNode("streamIntentRecognition", AsyncNodeAction.node_async(new StreamIntentRecognitionNode(chatClientBuilder, contextService)))
                // 流式症状详情收集节点  
                .addNode("streamSymptomDetail", AsyncNodeAction.node_async(new StreamSymptomDetailNode(chatClientBuilder, contextService)))
                // 科室推荐节点（可以保持原有的，也可以改为流式）
                .addNode("departmentRecommendation", AsyncNodeAction.node_async(new DepartmentRecommendationNode(departmentRecommendationService)))
                // 无意图问题处理节点
                .addNode("noIntentQuestion", AsyncNodeAction.node_async(new NoIntentQuestionNode(conversationHandlingService)))

                // 边：START -> 流式意图识别
                .addEdge(StateGraph.START, "streamIntentRecognition")

                // 条件边：流式意图识别 -> 根据意图类型分流
                .addConditionalEdges("streamIntentRecognition",
                        state -> {
                            String intentType = state.value("intentType", "无意图");
                            System.out.println("=== 流式意图识别结果: " + intentType + " ===");
                            if ("有意图".equals(intentType)) {
                                return CompletableFuture.completedFuture("streamSymptomDetail");
                            } else if ("直接挂号".equals(intentType)) {
                                System.out.println("=== 流式检测到直接挂号，跳转到科室推荐 ===");
                                return CompletableFuture.completedFuture("departmentRecommendation");
                            } else {
                                return CompletableFuture.completedFuture("noIntentQuestion");
                            }
                        },
                        Map.of(
                                "streamSymptomDetail", "streamSymptomDetail",
                                "departmentRecommendation", "departmentRecommendation",
                                "noIntentQuestion", "noIntentQuestion"
                        ))

                // 条件边：流式症状详情收集 -> 智能判断是否继续询问或推荐科室
                .addConditionalEdges("streamSymptomDetail",
                        state -> {
                            Boolean shouldContinue = state.value("shouldContinue", true);
                            Boolean readyForRecommendation = state.value("readyForRecommendation", false);
                            Integer questionCount = state.value("questionCount", 0);
                            Double completenessScore = state.value("completenessScore", 0.5);
                            String reasoning = state.value("reasoning", "继续收集");
                            
                            System.out.println("=== 流式症状收集状态 ===");
                            System.out.println("问题轮次: " + questionCount);
                            System.out.println("是否继续: " + shouldContinue);
                            System.out.println("准备推荐: " + readyForRecommendation);
                            System.out.println("完整性评分: " + completenessScore);
                            System.out.println("推理: " + reasoning);

                            // 智能判断：基于AI分析结果和完整性评分
                            if (readyForRecommendation || !shouldContinue || questionCount >= 4 || completenessScore >= 0.8) {
                                System.out.println("=== 触发科室推荐条件 ===");
                                return CompletableFuture.completedFuture("departmentRecommendation");
                            } else {
                                System.out.println("=== 等待用户下一轮输入 ===");
                                return CompletableFuture.completedFuture(StateGraph.END);
                            }
                        },
                        Map.of(
                                StateGraph.END, StateGraph.END,
                                "departmentRecommendation", "departmentRecommendation"
                        ))

                // 边：科室推荐 -> END
                .addEdge("departmentRecommendation", StateGraph.END)
                // 边：无意图问题 -> END
                .addEdge("noIntentQuestion", StateGraph.END);

        // 打印流式图结构
        GraphRepresentation representation = streamGraph.getGraph(GraphRepresentation.Type.PLANTUML,
                "流式智能问诊流程图");
        System.out.println("\n=== 流式智能问诊 UML Flow ===");
        System.out.println(representation.content());
        System.out.println("==================================\n");

        return streamGraph;
    }
}